package com.xyh.ollama_test.init;

import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections.VectorParams;
import io.qdrant.client.grpc.Collections.Distance;
import io.qdrant.client.grpc.Collections.CollectionOperationResponse;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.stereotype.Component;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;
import java.util.List;

@Component
public class QdrantInitializer implements ApplicationRunner {
    private final QdrantClient qdrantClient;
    private static final String COLLECTION_NAME = "text_embeddings";
    private static final int VECTOR_SIZE = 1024; // BGE-M3的向量维度，根据实际情况调整

    public QdrantInitializer(QdrantClient qdrantClient) {
        this.qdrantClient = qdrantClient;
    }

    @Override
    public void run(ApplicationArguments args) throws Exception {
        // 检查集合是否存在，不存在则创建
        ListenableFuture<List<String>> listCollectionsFuture = qdrantClient.listCollectionsAsync();

        Futures.addCallback(listCollectionsFuture, new FutureCallback<List<String>>() {
            @Override
            public void onSuccess(List<String> collections) {
                if (!collections.contains(COLLECTION_NAME)) {
                    ListenableFuture<CollectionOperationResponse> createCollectionFuture = qdrantClient.createCollectionAsync(
                            COLLECTION_NAME,
                            VectorParams.newBuilder()
                                    .setSize(VECTOR_SIZE)
                                    .setDistance(Distance.Cosine)
                                    .build()
                    );

                    Futures.addCallback(createCollectionFuture, new FutureCallback<CollectionOperationResponse>() {
                        @Override
                        public void onSuccess(CollectionOperationResponse result) {
                            System.out.println("Collection created successfully");
                        }

                        @Override
                        public void onFailure(Throwable t) {
                            System.err.println("Failed to create collection: " + t.getMessage());
                        }
                    }, MoreExecutors.directExecutor());
                }
            }

            @Override
            public void onFailure(Throwable t) {
                System.err.println("Failed to list collections: " + t.getMessage());
            }
        }, MoreExecutors.directExecutor());
    }
}